﻿/*
 * The MIT License (MIT)
 * 
 * Copyright (c) 2012 Brno University of Technology - Faculty of Information Technology (http://www.fit.vutbr.cz)
 * Author(s):
 * Vladimir Vesely (mailto:ivesely@fit.vutbr.cz)
 * Martin Mares
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated 
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation 
 * the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, 
 * and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED 
 * TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, 
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

using System;
using System.Collections.Generic;
using System.IO;
using CommandLine;
using CommandLine.Text;


namespace PcapMerger
{
    class Program
    {

        private static readonly HeadingInfo HeadingInfo = new HeadingInfo("PCAPMerger 2.1 by Vladimir Vesely (ivesely@fit.vutbr.cz) ");
        private class Options
        {
            [Option(null, "mnm1", DefaultValue = false, Required = false, HelpText = "Print information in MNM Capture File Header")]
            public Boolean MNM1 { get; set; }
            [Option(null, "mnm2", DefaultValue = false, Required = false, HelpText = "Print information in MNM Process Info Table")]
            public Boolean MNM2 { get; set; }
            [Option(null, "mnm3", DefaultValue = false, Required = false, HelpText = "Print information in MNM Extended Info")]
            public Boolean MNM3 { get; set; }
            [Option(null, "mnm4", DefaultValue = false, Required = false, HelpText = "Print information in MNM Frame Table Layout")]
            public Boolean MNM4 { get; set; }

            [Option(null, "tcpd1", DefaultValue = false, Required = false, HelpText = "Print information in TCPD Global Header")]
            public Boolean TCPD1 { get; set; }
            [Option(null, "tcpd2", DefaultValue = false, Required = false, HelpText = "Print information in TCPD Frame Table")]
            public Boolean TCPD2 { get; set; }

            [OptionList("i", "input", Separator = '+', Required = true, HelpText = "Input PCAP files to be processed delimited with + sign.")]
            public IList<string> InputFiles { get; set; }

            [Option("o", "output", DefaultValue = "out.pcap", Required = true, HelpText = "An output PCAP file")]
            public String OutputFileName { get; set; }
            [Option("f", "format", DefaultValue = "tcpd", Required = true, HelpText = "Output PCAP format, currently supports 'tcpd' (standard libPCAP), 'mnm' (MS NetMon) and 'pcapng' (NextGeneration PCAP).")]
            public String OutputFormat { get; set; }

            [Option("v", "verbose", DefaultValue = false, Required = false, HelpText = "Verbose mode with additional information. All debug and informational messages will be displayed!")]
            public Boolean Verbose { get; set; }
            [Option("s", "sorting", DefaultValue = "quicksort", Required = false, HelpText = "Used sorting alghoritm, currently you can choose between 'quicksort' or 'heapsort'")]
            public String SortAlg { get; set; }


            [HelpOption(HelpText = "Display this help screen.")]
            public string GetUsage()
            {
                var help = new HelpText(Program.HeadingInfo);
                help.AdditionalNewLineAfterOption = false;
                help.Copyright = new CopyrightInfo("NES@FIT - http://www.fit.vutbr.cz/research/groups/nes@fit/ \nBrno University of Technology, Faculty of Information Technology", 2012);
                help.AddPreOptionsLine("Usage: PCAPMerger.exe -i file1.cap+file2.cap -o output.cap -s heapsort -f mnm");
                help.AddPreOptionsLine("Options:");
                help.AddOptions(this);
                help.AddPostOptionsLine("This program reads input PCAP files and merge their content sorted by time to output file regardless to any API.");
                return help;
            }
        }

        private static readonly Options Opt = new Options();


        private enum ExitCodes
        {
            Success = 0,
            ArgParse = 1,
            VerifyMediaTypes = 2,
            FileCheck = 3,
            WritingOut = 4
        }
        private static void EndPrompt(int exit)
        {
            PmConsolePrinter.PrintInfoEol("\nPress Enter to continue...");
            Console.ReadLine();
            Environment.Exit(exit);
        }

        static void Main(string[] args)
        {
            /* 1) Check for proper commandline arguments present */
            ParseArgs(args);

            var sfsList = new List<IPmCapture>();

            /* 2) Check input PCAP files existence and open them for binary read operations */
            CheckAndOpenSfsFiles(sfsList);

            /* 3) Get basic information from sniff files */
            ParseSfsFiles(sfsList);

            /* 3.5) Produce information to Console if necessary */
            foreach (var sfs in sfsList)
            {
                PmSupportedTypes.CaptureFileTypes type = sfs.FileType();

                if (type == PmSupportedTypes.CaptureFileTypes.MicrosoftNetworkMonitor && Opt.MNM1)
                    ((PmCaptureMNM)sfs).PrintMNMCaptureHeader();
                if (type == PmSupportedTypes.CaptureFileTypes.MicrosoftNetworkMonitor && ((PmCaptureMNM)sfs).HasProcessInfoTable() && (Opt.MNM2))
                    ((PmCaptureMNM)sfs).PrintMNMProcessInfoTable();
                if (type == PmSupportedTypes.CaptureFileTypes.MicrosoftNetworkMonitor && ((PmCaptureMNM)sfs).HasExtendedInfo() && (Opt.MNM3))
                    ((PmCaptureMNM)sfs).PrintMNMExtendedInfo();
                if (type == PmSupportedTypes.CaptureFileTypes.MicrosoftNetworkMonitor && (Opt.MNM4))
                    ((PmCaptureMNM)sfs).PrintMNMFrameTable();
                if (type == PmSupportedTypes.CaptureFileTypes.LibPCAP && (Opt.TCPD1))
                    ((PmCapturePcap)sfs).PrintTCPDGlobalHeader();
                if (type == PmSupportedTypes.CaptureFileTypes.LibPCAP && (Opt.TCPD2))
                    ((PmCapturePcap)sfs).PrintTCPDFrameTable();
            }

            /* 4) Verify same Link Layer Type, the one from the first file is used as reference */
           // VerifyMediaTypes(sfsList);

            /* 5) Sort frames from all files according to timestamp */
            var frames = SortFrames(sfsList);

            /* 6) Creates resulting LibPCAP file for binary writing */
            CreateOutput(frames, Opt.OutputFormat);

            /* 7) Close opened inputPCAPs */
            CloseFiles(sfsList);

            if (!Opt.Verbose) return;
            EndPrompt((int)ExitCodes.Success);

            if (!Opt.Verbose) return;
            PmConsolePrinter.PrintInfoEol("\nPress Enter to continue...");
            Console.ReadLine();

        }

        private static void CloseFiles(IEnumerable<IPmCapture> sfsList)
        {
            foreach (var sfs in sfsList)
                sfs.Close();
        }

        private static void CreateOutput(List<IPmFrame> frames, String format)
        {
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfo("Writing output file " + Opt.OutputFileName + " ... ");
            /* Produce output */

            IPmCapture output;

            switch (format.ToLower())
            {
                case "pcapng":
                    output = PmCaptureManager.CreateCaptureFile(PmSupportedTypes.CaptureFileTypes.PCAPNextGen);
                    break;
                case "mnm":
                    output = PmCaptureManager.CreateCaptureFile(PmSupportedTypes.CaptureFileTypes.MicrosoftNetworkMonitor);
                    break;
                case "tcpd":
                default:
                    output = PmCaptureManager.CreateCaptureFile(PmSupportedTypes.CaptureFileTypes.LibPCAP);
                    break;
            }

            if (output != null)
            {
                output.SetFrames(frames);
                if (!output.CreateOutput(Opt.OutputFileName))
                {
                    PmConsolePrinter.PrintError("Result file " + Opt.OutputFileName +
                                              " cannot be generated due to previous exception(s)!\n");
                    EndPrompt((int)ExitCodes.WritingOut);
                }
                else if (Opt.Verbose)
                {
                    PmConsolePrinter.PrintInfoEol("OK");
                }
            }

        }

        private static List<IPmFrame> SortFrames(IEnumerable<IPmCapture> sfsList)
        {
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfo("Sorting and merging frames ... ");
            var frames = PmCaptureManager.SortFramesIntoFV(sfsList, Opt.SortAlg);
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfoEol("OK");
            return frames;
        }

        //private static void InitializeFrames(List<RawParsing.PmPcapFile> sfsList)
        //{
        //    if (Opt.Verbose)
        //        Console.Write("Creating frame vectors ... ");
        //    /* Create PmFrameVector field by parsing information from each sfs */
        //    if (!RawParsing.CreateFrameVectors(sfsList))
        //    {
        //        RawParsing.PrintError("Creating of FrameVectors failed.");
        //        EndPrompt((int) ExitCodes.FrameVect);
        //    }
        //    else if (Opt.Verbose) Console.WriteLine("OK");
        //}


        private static void ParseSfsFiles(IEnumerable<IPmCapture> sfsList)
        {
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfo("Parsing headers of PCAP files ... ");
            foreach (var sfs in sfsList)
            {
                sfs.CreateFrameTable();
            }
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfoEol("OK");
        }

       /* private static void VerifyMediaTypes(List<IPmCapture> sfsList)
        {
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfo("Verifying used media types integrity ... ");
            if (!PmCaptureManager.VerifySfsMediaTypes(sfsList))
            {
                PmConsolePrinter.PrintError("VerifyMediaTypes failed. No files left for parsing!");
                EndPrompt((int)ExitCodes.VerifyMediaTypes);
            }
            else if (Opt.Verbose) PmConsolePrinter.PrintInfoEol("OK");
        }*/

        private static void CheckAndOpenSfsFiles(ICollection<IPmCapture> sfsList)
        {
            if (Opt.Verbose)
                PmConsolePrinter.PrintInfoEol("Input file(s): ");
            foreach (var ff in Opt.InputFiles)
            {
                /* Open input files */
                if (Opt.Verbose)
                    PmConsolePrinter.PrintInfo(ff + " ... ");

                if (File.Exists(Path.GetFullPath(ff)))
                {
                    try
                    {                    
                        sfsList.Add(PmCaptureManager.CreateCaptureFile(ff));
                    }
                    catch (Exception ex)
                    {
                        PmConsolePrinter.PrintError("Error>\n" + ex.Message);
                        PmConsolePrinter.PrintError("\nHence, skipping processing of file " + ff);
                        //throw;
                    }
                    if (Opt.Verbose)
                        PmConsolePrinter.PrintInfoEol("OK");
                }
                else
                {
                    PmConsolePrinter.PrintError("File " + ff + " does not exist!");
                    //EndPrompt((int)ExitCodes.FileCheck);
                }



            }
        }

        private static void ParseArgs(string[] args)
        {
            ICommandLineParser cmdlnParser = new CommandLineParser(new CommandLineParserSettings(Console.Error));
            if (!cmdlnParser.ParseArguments(args, Opt))
                EndPrompt((int)ExitCodes.ArgParse);
        }
    }
}


